Open In Colab

How to Train YOLOv5 on Custom Objects

This tutorial is based on the YOLOv5 repository by Ultralytics. This notebook shows training on your own custom objects. Many thanks to Ultralytics for putting this repository together - we hope that in combination with clean data management tools at Roboflow, this technologoy will become easily accessible to any developer wishing to use computer vision in their projects.

Accompanying Blog Post

We recommend that you follow along in this notebook while reading the blog post on how to train YOLOv5, concurrently.

Steps Covered in this Tutorial

In this tutorial, we will walk through the steps required to train YOLOv5 on your custom objects. We use a public blood cell detection dataset, which is open source and free to use. You can also use this notebook on your own data.

To train our detector we take the following steps:

About

Roboflow enables teams to deploy custom computer vision models quickly and accurately. Convert data from to annotation format, assess dataset health, preprocess, augment, and more. It's free for your first 1000 source images.

Looking for a vision model available via API without hassle? Try Roboflow Train.

Roboflow Wordmark

Install Dependencies

(Remember to choose GPU in Runtime if not already selected. Runtime --> Change Runtime Type --> Hardware accelerator --> GPU)

Define Model Configuration and Architecture

We will write a yaml script that defines the parameters for our model like the number of classes, anchors, and each layer.

You do not need to edit these cells, but you may.

Train Custom YOLOv5 Detector

Next, we'll fire off training!

Here, we are able to pass a number of arguments:

Evaluate Custom YOLOv5 Detector Performance

Training losses and performance metrics are saved to Tensorboard and also to a logfile defined above with the --name flag when we train. In our case, we named this yolov5s_results. (If given no name, it defaults to results.txt.) The results file is plotted as a png after training completes.

Note from Glenn: Partially completed results.txt files can be plotted with from utils.utils import plot_results; plot_results().

Curious? Visualize Our Training Data with Labels

After training starts, view train*.jpg images to see training images, labels and augmentation effects.

Note a mosaic dataloader is used for training (shown below), a new dataloading concept developed by Glenn Jocher and first featured in YOLOv4.

Run Inference With Trained Weights

Run inference with a pretrained checkpoint on contents of test/images folder downloaded from Roboflow.